import torch
import torchvision

class classifier(torch.nn.Module):
    '''
    kwargs:
    {
        'classes' : 类别数,
        'torch_pretrain' : 是否使用torchvision的预训练权重,
    }
    '''
    def __init__(self, opt):
        super().__init__()

        classes = opt.classes
        torch_pretrain = opt.torch_pretrain

        self.network = torchvision.models.resnet101(pretrained=torch_pretrain)
        num_features = self.network.fc.in_features
        self.network.fc = torch.nn.Linear(num_features, classes)

        # print(self.network)

    def forward(self, input):
        x = input
        output = self.network(x)

        return output